package com.leetcode;

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 347. 前 K 个高频元素
 * 通过堆排序 维护一个n - k个元素的最大堆 时间复杂度O(nlog(n-k))
 */
public class Solution347_2 {

    public int[] topKFrequent(int[] nums, int k) {
        int n = nums.length;
        Map<Integer, Integer> map = new HashMap<>(n);
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }
        // java默认的堆是最小堆
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>((Integer a, Integer b) -> map.get(b) - map.get(a));
        for (Integer num : map.keySet()) {
            if (priorityQueue.size() >= (map.size() - k)) {
                // 这里有一个小的陷阱， 如果堆为空则肯定是不能去poll和peek操作的
                if (!priorityQueue.isEmpty() && map.get(num) < map.get(priorityQueue.peek())) {
                    priorityQueue.poll();
                    priorityQueue.offer(num);
                }
            } else {
                priorityQueue.offer(num);
            }
        }
        while (!priorityQueue.isEmpty()) {
            map.remove(priorityQueue.poll());
        }
        return map.keySet().stream().mapToInt(a -> a).toArray();
    }

    public static void main(String[] args) {
//        int[] nums = new int[]{1, 1, 1, 2, 2, 3};
//        int k = 2;

//        int[] nums = new int[]{1};
//        int k = 1;

        int[] nums = new int[]{3, 0, 1, 0};
        int k = 1;
        int[] res = new Solution347_2().topKFrequent(nums, k);
        for (int num : res) {
            System.out.println(num + " ");
        }
        System.out.println();
    }

}
